import torch
import torchvision
import torch.nn.functional as fun
from torch import nn



class Discriminator_X(nn.Module): # the discriminator of set x

    def __init__(self, *args, **kwargs) -> None:
        super(Discriminator_X, self).__init__(*args, **kwargs)

    def forward():
        pass


class Discriminator_Y(nn.Module): # the discriminator of set y

    def __init__(self, *args, **kwargs) -> None:
        super(Discriminator_Y, self).__init__(*args, **kwargs)

    def forward():
        pass


class Discriminator(nn.Module):

    def __init__(self, *args, **kwargs) -> None:
        super(Discriminator, self).__init__(*args, **kwargs)
        self.main = nn.Sequential(*[
            nn.Conv2d(3, 64, 4, 2, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, 4, 2, 1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Conv2d(256, 512, 4, 1, 1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2),
            nn.Conv2d(512, 1, 4, 1, 1),
        ])

    def forward(self, x):
        return self.main(x)
    
